///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/bfgs.cpp
///
/// \brief   implementation file for the bfgs updater
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<vector>
#include<numeric>
#include<cassert>
#include<algorithm>
#include<cmath>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/fqmc/bfgs.h>
#include<formic/wfn/configuration.h>
#include<formic/input/base.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/random/random.h>
#include<formic/output/output.h>
#include<formic/exception.h>
#include<formic/lapack/interface.h>
#include<formic/timing/timing.h>
#include<formic/fqmc/line_search.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  constructs the BFGS updater
///
/// \param[in]      userinp        the user's input options
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> formic::fqmc::BFGSUpdater<S>::BFGSUpdater(const formic::InputBase & userinp)
{

}

template formic::fqmc::BFGSUpdater<double>::BFGSUpdater(const formic::InputBase & userinp);
template formic::fqmc::BFGSUpdater<std::complex<double> >::BFGSUpdater(const formic::InputBase & userinp);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  performs a bfgs update on the supplied wave function
///
/// \param[in]      userinp        the user's input options
/// \param[in]      pmtr_hist      the history of the parameters
/// \param[in]      grad_hist      the history of the gradients
/// \param[in]      sampler        object that controls sampling
/// \param[in,out]  co             a shared pointer to the value function
///                                On exit, the best update has been added to the value function
///
/// \return Whether or not the update was rejected.  Updates are rejected if they do not
///         lower the energy.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> bool formic::fqmc::BFGSUpdater<S>::update(const formic::InputBase & userinp,
                                                             const std::vector<boost::shared_ptr<std::vector<double> > > & pmtr_hist,
                                                             const std::vector<boost::shared_ptr<std::vector<double> > > & grad_hist,
                                                             boost::shared_ptr<formic::fqmc::Sampler<S> > & sampler,
                                                             const boost::shared_ptr<formic::CoeffObj<S> > & co) {

  struct local_funcs {

    static void real_to_cplx(const size_t n, const double * const r, double * const c) {
      std::copy(r, r+n, c);
    }

    static void real_to_cplx(const size_t n, const double * const r, std::complex<double> * const c) {
      for (size_t i = 0; i < n; i++)
        c[i].real() = r[i];
      for (size_t i = 0; i < n; i++)
        c[i].imag() = r[i+n];
    }

  };

  //// print a greeting
  //if (formic::mpi::rank() == 0)
  //  formic::of << "performing a BFGS update." << std::endl << std::endl;

  // ensure we are using exact sampling (bfgs code does not currently accomodate stochastic sampling)
  if ( !userinp.get<bool>("exact_sampling") )
    throw formic::Exception("bfgs update only implemented for exact sampling");

  std::vector<double> pdiff(grad_hist.at(0)->size(), 0.0);
  std::vector<double> gdiff(grad_hist.at(0)->size(), 0.0);

  // compute the rho vector
  std::vector<double> rho(pmtr_hist.size(), 0.0);
  for (int i = 1; i < pmtr_hist.size() && i <= userinp.get<int>("bfgs_hist_len"); i++) {
    //rho.at(i) += formic::xdotc(pmtr_hist.at(i-1)->size(), &pmtr_hist.at(i-1)->at(0), 1, &grad_hist.at(i-1)->at(0), 1);
    //rho.at(i) -= formic::xdotc(pmtr_hist.at(i-1)->size(), &pmtr_hist.at(i-0)->at(0), 1, &grad_hist.at(i-1)->at(0), 1);
    //rho.at(i) -= formic::xdotc(pmtr_hist.at(i-1)->size(), &pmtr_hist.at(i-1)->at(0), 1, &grad_hist.at(i-0)->at(0), 1);
    //rho.at(i) += formic::xdotc(pmtr_hist.at(i-1)->size(), &pmtr_hist.at(i-0)->at(0), 1, &grad_hist.at(i-0)->at(0), 1);
    //rho.at(i) = 1.0 / rho.at(i);
    formic::xcopy(pdiff.size(),       &pmtr_hist.at(i-1)->at(0), 1, &pdiff.at(0), 1);
    formic::xaxpy(pdiff.size(), -1.0, &pmtr_hist.at(i-0)->at(0), 1, &pdiff.at(0), 1);
    formic::xcopy(gdiff.size(),       &grad_hist.at(i-1)->at(0), 1, &gdiff.at(0), 1);
    formic::xaxpy(gdiff.size(), -1.0, &grad_hist.at(i-0)->at(0), 1, &gdiff.at(0), 1);
    rho.at(i) = 1.0 / formic::xdotc(pdiff.size(), &pdiff.at(0), 1, &gdiff.at(0), 1);
    //if (formic::mpi::rank() == 0)
    //  formic::of << boost::format("rho[%2i] = %.2e") % i % rho.at(i) << std::endl;
  }

  // prepare space for the alpha and beta vectors
  std::vector<double> alpha(pmtr_hist.size(), 0.0);
  std::vector<double> beta(pmtr_hist.size(), 0.0);

  // compute the update
  std::vector<double> update(**grad_hist.begin());
  for (int i = 1; i < pmtr_hist.size() && i <= userinp.get<int>("bfgs_hist_len"); i++) {
    //alpha.at(i) += rho.at(i) * formic::xdotc(update.size(), &pmtr_hist.at(i-1)->at(0), 1, &update.at(0), 1);
    //alpha.at(i) -= rho.at(i) * formic::xdotc(update.size(), &pmtr_hist.at(i-0)->at(0), 1, &update.at(0), 1);
    //formic::xaxpy(update.size(), -alpha.at(i), &grad_hist.at(i-1)->at(0), 1, &update.at(0), 1);
    //formic::xaxpy(update.size(),  alpha.at(i), &grad_hist.at(i-0)->at(0), 1, &update.at(0), 1);
    formic::xcopy(pdiff.size(),       &pmtr_hist.at(i-1)->at(0), 1, &pdiff.at(0), 1);
    formic::xaxpy(pdiff.size(), -1.0, &pmtr_hist.at(i-0)->at(0), 1, &pdiff.at(0), 1);
    formic::xcopy(gdiff.size(),       &grad_hist.at(i-1)->at(0), 1, &gdiff.at(0), 1);
    formic::xaxpy(gdiff.size(), -1.0, &grad_hist.at(i-0)->at(0), 1, &gdiff.at(0), 1);
    alpha.at(i) = rho.at(i) * formic::xdotc(update.size(), &pdiff.at(0), 1, &update.at(0), 1);
    formic::xaxpy(update.size(), -alpha.at(i), &gdiff.at(0), 1, &update.at(0), 1);
  }
  for (int i = std::min(int(pmtr_hist.size() - 1), userinp.get<int>("bfgs_hist_len")); i > 0; i--) {
    //beta.at(i) += rho.at(i) * formic::xdotc(update.size(), &grad_hist.at(i-1)->at(0), 1, &update.at(0), 1);
    //beta.at(i) -= rho.at(i) * formic::xdotc(update.size(), &grad_hist.at(i-0)->at(0), 1, &update.at(0), 1);
    //formic::xaxpy(update.size(), alpha.at(i) - beta.at(i), &pmtr_hist.at(i-1)->at(0), 1, &update.at(0), 1);
    //formic::xaxpy(update.size(), beta.at(i) - alpha.at(i), &pmtr_hist.at(i-0)->at(0), 1, &update.at(0), 1);
    formic::xcopy(pdiff.size(),       &pmtr_hist.at(i-1)->at(0), 1, &pdiff.at(0), 1);
    formic::xaxpy(pdiff.size(), -1.0, &pmtr_hist.at(i-0)->at(0), 1, &pdiff.at(0), 1);
    formic::xcopy(gdiff.size(),       &grad_hist.at(i-1)->at(0), 1, &gdiff.at(0), 1);
    formic::xaxpy(gdiff.size(), -1.0, &grad_hist.at(i-0)->at(0), 1, &gdiff.at(0), 1);
    beta.at(i) = rho.at(i) * formic::xdotc(update.size(), &gdiff.at(0), 1, &update.at(0), 1);
    formic::xaxpy(update.size(), alpha.at(i)-beta.at(i), &pdiff.at(0), 1, &update.at(0), 1);
  }
  formic::xscal(update.size(), -1.0, &update.at(0), 1);

  // print the dot product of the gradient and the update direction (this should be negative)
  if (formic::mpi::rank() == 0)
    formic::of << boost::format("dot of gradient with BFGS search direction = %20.6e")
                  % formic::xdotc(update.size(), &update.at(0), 1, &(**grad_hist.begin()).at(0), 1)
               << std::endl
               << std::endl;

  // convert the update to complex form if necessary
  std::vector<S> cplx_update(co->n_der(), formic::zero(S()));
  local_funcs::real_to_cplx(co->n_der(), &update.at(0), &cplx_update.at(0));

  // perform a line search along the update direction and remember the optimum step size
  double step_arg = userinp.get<double>("opt_bfgs_step_size");
  formic::fqmc::line_search(userinp, &cplx_update.at(0), userinp.get<double>("update_max_sdev"), sampler, co, step_arg);

  // return whether the update was rejected
  bool rejected = ( step_arg == 0.0 );
  return rejected;

}

template bool formic::fqmc::BFGSUpdater<double>::update(const formic::InputBase & userinp,
                                                        const std::vector<boost::shared_ptr<std::vector<double> > > & pmtr_hist,
                                                        const std::vector<boost::shared_ptr<std::vector<double> > > & grad_hist,
                                                        boost::shared_ptr<formic::fqmc::Sampler<double> > & sampler,
                                                        const boost::shared_ptr<formic::CoeffObj<double> > & co);
template bool formic::fqmc::BFGSUpdater<std::complex<double> >::update(const formic::InputBase & userinp,
                                                                       const std::vector<boost::shared_ptr<std::vector<double> > > & pmtr_hist,
                                                                       const std::vector<boost::shared_ptr<std::vector<double> > > & grad_hist,
                                                                       boost::shared_ptr<formic::fqmc::Sampler<std::complex<double> > > & sampler,
                                                                       const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & co);
